from queue import Queue
import sys
import os

# 获取当前文件所在目录的绝对路径
current_dir = os.path.dirname(os.path.abspath(__file__))
# 添加项目根目录到Python路径
sys.path.append(os.path.join(current_dir, '../'))  # 根据实际位置调整层级
from concurrent.futures import ThreadPoolExecutor, as_completed
from py_utils.rknn_executor import RKNN_model_container 

def initRKNN(rknnModel="./models/merge_yolov11_seg.rknn", id=0):
    rknn_lite = RKNN_model_container(rknnModel, "rk3588", id)
    return rknn_lite


def initRKNNs(rknnModel="./models/merge_yolov11_seg.rknn", TPEs=1,npus=[0]):
    rknn_list = []
    for npu in npus:
        for i in range(TPEs):
            rknn_list.append(initRKNN(rknnModel, npu))
    return rknn_list


class rknnPoolExecutor():
    def __init__(self, rknnModel, TPEs, func,npus=[0]):
        self.TPEs = TPEs
        self.queue = Queue()
        self.rknnPool = initRKNNs(rknnModel, TPEs,npus)
        self.pool = ThreadPoolExecutor(max_workers=TPEs)
        self.func = func
        self.num = 0

    def put(self, frame):
        self.queue.put(self.pool.submit(
            self.func, self.rknnPool[self.num % self.TPEs], frame))
        self.num += 1

    def get(self):
        if self.queue.empty():
            return None, False
        fut = self.queue.get()
        return fut.result(), True

    def release(self):
        self.pool.shutdown()
        for rknn_lite in self.rknnPool:
            rknn_lite.release()